package kafka

import java.sql.{Connection, DriverManager, PreparedStatement}

import org.apache.spark.sql.{DataFrame, Dataset, ForeachWriter, Row, SparkSession}

object TestStreamingKafkaMysql {
  def main(args: Array[String]): Unit = {

    // 准备环境
    val spark: SparkSession = SparkSession.builder()
      .appName("test01")
      .master("local[*]")
      .getOrCreate()


    // 设置日志级别
    spark.sparkContext.setLogLevel("WARN")

    //  导入隐式转换
    import spark.implicits._

    val kafkaDatas: DataFrame = spark.readStream.format("kafka")
      .option("kafka.bootstrap.servers", "dw-Huake05:9092,dw-Huake06:9092,dw-Huake07:9092")
      .option("subscribe", "test01")
      .load()

    //3.处理数据
    //注意:StructuredStreaming整合Kafka获取到的数据都是字节类型,所以需要按照官网要求,转成自己的实际类型
    // kafkaDatas  内部数据是kafka数据(key,value)
    val kafkaDataString: Dataset[(String, String)] = kafkaDatas.selectExpr("CAST(key AS string)", "CAST(value AS string)").as[(String, String)]

    // 处理
    val word: Dataset[String] = kafkaDataString.flatMap(x => x._2.split(" "))
    val wordCount: Dataset[Row] = word.groupBy("value").count().sort($"count".desc)


    val intoMYSQL: intoMysql = new intoMysql("jdbc:mysql://dw-Huake05:3306/testspark?characterEncoding=UTF-8", "root", "huake@2021")


    // 输出
    wordCount.writeStream.foreach(intoMYSQL)   // 遍历流中的数据,调用功能类,将其保存/更新到数据库
      .outputMode("complete")   // 设置输出模式
      .start()  // 开启任务
      .awaitTermination()   // 等待程序结束


  }

  // 创建一个类，编写将数据更新/插入到mysql数据库的代码
  class intoMysql(url: String, username: String, password: String) extends ForeachWriter[Row] with Serializable {

    // 准备连接对象
    var connection: Connection = _
    // 设置sql
    var preparedStatement: PreparedStatement = _


    // 用于打开数据库连接
    override def open(partitionId: Long, version: Long): Boolean = {

      // 获取链接
      connection = DriverManager.getConnection(url, username, password)
      //获取链接无误返回True
      true
    }

    // 用于更新/插入数据到mysql
    override def process(value: Row): Unit = {
      // value内的第一个数是单词
      val word: String = value.get(0).toString

      // value内的第二个数是单词的数量
      val count: Int = value.get(1).toString.toInt

      // 打印结果
      println("word:" + word + "\tcount:" + count)

      //REPLACE INTO:表示如果表中没有数据这插入,如果有数据则替换
      //注意:REPLACE INTO要求表有主键或唯一索引
      val sql = "REPLACE INTO `t_word` (`id`, `word`, `count`) VALUES (NULL, ?, ?)"
      // val sql = "INSERT INTO t_word (id, word, count) VALUES (null,?,?) ON DUPLICATE KEY UPDATE count = ?"

      val prepareStatement: PreparedStatement = connection.prepareStatement(sql)
      // 设置字段
      prepareStatement.setString(1, word)
      prepareStatement.setInt(2, count)
      //prepareStatement.setInt(3, count)

      // 执行
      prepareStatement.executeUpdate()

    }

    // 关闭数据库连接
    override def close(errorOrNull: Throwable): Unit = {


      if (connection != null) {
        connection.close()
      }
      if (preparedStatement != null) {
        preparedStatement.close()
      }
    }
  }
}

